#!/usr/bin/env python3
# Author: Armit
# Create Time: 2023/01/09 

# 一个 structure 下的所有 model 的 seq 必然一致吗
# 换句话是，一个 structure 是一个物理客体，一个 model 是对它的一次拍照（？
# => 

from preprocess_pdb import *


def worker(tasks:List[Path], pdb_fps:List[Path]):
  parser = PDBParser()

  for file in tasks:
    try:
      with gzip.open(file, 'rt', encoding='utf-8') as fh:
        id = file.stem[:-len('.pdb')]
        structure = parser.get_structure(id, fh)
    except:
      print(f'<< error ignore file {file}')
      continue

    models = []
    for model in structure.get_models():
      chains = []
      for chain in model.get_chains():
        chains.append(''.join([r.resname for r in chain.get_residues()]))
      models.append(chains)

    if len(models) > 1:
      model0 = models[0]
      for i in range(1, len(models)):
        if model0 != models[i]:
          pdb_fps.append(str(file))
          break


def test_pdb_seq_uniq():
  # multiprocess tasks & workers
  tasks: List[Path] = list(walk_for_pdb())
  n_tasks_each = math.ceil(len(tasks) / N_WORKER)
  print(f'>> found {len(tasks)} pdb files')
  print(f'      assign {n_tasks_each} tasks for each thread')

  manager = Manager()
  pdb_fps: List[str] = manager.list()
  procs = [
    Process(target=worker, 
            args=(tasks[i*n_tasks_each:(i+1)*n_tasks_each], pdb_fps), 
            daemon=True) 
      for i in range(N_WORKER)
  ]
  for proc in procs: proc.start()
  for proc in procs: proc.join()
  del procs, manager, tasks

  pdb_fps = list(set(pdb_fps))
  print(f'>> found {len(pdb_fps)} items')

  with open(Path(DATA_PATH) / 'test_pdb_seq_uniq.txt', 'w', encoding='utf-8') as fh:
    for fp in pdb_fps:
      fh.write(fp)
      fh.write('\n')


if __name__ == '__main__':
  test_pdb_seq_uniq()
